import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QPushButton, QLabel, QFileDialog, 
                             QMessageBox)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from PIL import Image, ImageEnhance

class ImageSegmentationGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.model = None
        self.original_image = None
        self.segmented_image = None
        
        # 尝试加载模型
        self.load_model()
    
    def init_ui(self):
        self.setWindowTitle('图像分割工具')
        self.setGeometry(100, 100, 1200, 600)
        
        # 创建主窗口部件
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        
        # 创建布局
        main_layout = QHBoxLayout()
        main_widget.setLayout(main_layout)
        
        # 左侧控制面板
        control_layout = QVBoxLayout()
        control_widget = QWidget()
        control_widget.setLayout(control_layout)
        control_widget.setFixedWidth(200)
        
        # 选择图片按钮
        self.select_button = QPushButton('选择图片')
        self.select_button.clicked.connect(self.select_image)
        control_layout.addWidget(self.select_button)
        
        # 分割按钮
        self.segment_button = QPushButton('执行分割')
        self.segment_button.clicked.connect(self.perform_segmentation)
        self.segment_button.setEnabled(False)
        control_layout.addWidget(self.segment_button)
        
        # 加载模型按钮
        self.load_model_button = QPushButton('加载模型')
        self.load_model_button.clicked.connect(self.load_model)
        control_layout.addWidget(self.load_model_button)
        
        # 状态标签
        self.status_label = QLabel('请选择一张图片')
        self.status_label.setWordWrap(True)
        control_layout.addWidget(self.status_label)
        
        control_layout.addStretch()
        
        # 右侧图像显示区域
        image_layout = QHBoxLayout()
        
        # 原始图像显示
        self.original_image_label = QLabel('原始图像')
        self.original_image_label.setAlignment(Qt.AlignCenter)
        self.original_image_label.setMinimumSize(400, 400)
        self.original_image_label.setStyleSheet("border: 1px solid black;")
        image_layout.addWidget(self.original_image_label)
        
        # 分割结果图像显示
        self.segmented_image_label = QLabel('分割结果')
        self.segmented_image_label.setAlignment(Qt.AlignCenter)
        self.segmented_image_label.setMinimumSize(400, 400)
        self.segmented_image_label.setStyleSheet("border: 1px solid black;")
        image_layout.addWidget(self.segmented_image_label)
        
        # 添加到主布局
        main_layout.addWidget(control_widget)
        main_layout.addLayout(image_layout)
    
    def load_model(self):
        """加载训练好的模型"""
        try:
            model_path = 'model.h5'
            if os.path.exists(model_path):
                # 禁用GPU以避免兼容性问题
                os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
                self.model = tf.keras.models.load_model(model_path)
                self.status_label.setText('模型加载成功')
                self.segment_button.setEnabled(True)
            else:
                self.status_label.setText('未找到模型文件，请确保model.h5在当前目录')
                self.segment_button.setEnabled(False)
        except Exception as e:
            self.status_label.setText(f'模型加载失败: {str(e)}')
            self.segment_button.setEnabled(False)
    
    def select_image(self):
        """选择图片文件"""
        file_name, _ = QFileDialog.getOpenFileName(
            self, '选择图片', '', 'Images (*.png *.jpg *.jpeg)')
        
        if file_name:
            try:
                # 加载并显示原始图像
                self.original_image = load_img(file_name)
                self.display_image(self.original_image, self.original_image_label)
                
                # 清除之前的分割结果
                self.segmented_image_label.clear()
                self.segmented_image_label.setText('分割结果')
                self.segmented_image = None
                
                self.status_label.setText(f'已选择图片: {os.path.basename(file_name)}')
                self.segment_button.setEnabled(self.model is not None)
            except Exception as e:
                QMessageBox.critical(self, '错误', f'无法加载图片: {str(e)}')
    
    def display_image(self, image, label):
        """在指定标签中显示图片"""
        # 调整图像大小以适应显示区域
        image = image.copy()  # 创建副本避免修改原图
        image.thumbnail((400, 400), Image.LANCZOS)
        
        # 转换为QPixmap并显示
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        data = image.tobytes("raw", "RGB")
        qimage = QImage(data, image.size[0], image.size[1], QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qimage)
        label.setPixmap(pixmap)
    
    def perform_segmentation(self):
        """执行图像分割"""
        if self.original_image is None or self.model is None:
            return
        
        try:
            self.status_label.setText('正在分割图像...')
            QApplication.processEvents()  # 更新界面
            
            # 预处理图像 - 调整大小以匹配模型输入要求
            target_size = (160, 160)  # 与训练时的尺寸一致
            img = self.original_image.resize(target_size)
            
            # 确保图像是RGB格式（3通道）
            if img.mode != 'RGB':
                img = img.convert('RGB')
            
            # 转换为数组并调整维度以匹配模型期望的3通道输入
            img_array = img_to_array(img)  # 这会产生 (160, 160, 3) 的形状
            img_array = np.expand_dims(img_array, axis=0)  # 添加批次维度 (1, 160, 160, 3)
            img_array /= 255.0  # 归一化
            
            # 执行预测
            prediction = self.model.predict(img_array)
            
            # 后处理预测结果
            segmented = np.argmax(prediction[0], axis=-1)
            segmented_img = self.visualize_segmentation(segmented)
            
            # 转换为PIL图像并显示
            self.segmented_image = segmented_img
            self.display_image(self.segmented_image, self.segmented_image_label)
            
            self.status_label.setText('分割完成')
        except Exception as e:
            QMessageBox.critical(self, '错误', f'分割失败: {str(e)}')
            self.status_label.setText('分割失败')
    
    def visualize_segmentation(self, mask):
        """
        可视化分割结果为灰度图像
        :param mask: 分割掩码
        :return: 灰度分割结果图像
        """
        # 创建灰度图像表示分割结果
        # 将类别标签映射到不同的灰度值
        # 类别0 -> 0 (黑色)
        # 类别1 -> 85 (深灰色)
        # 类别2 -> 170 (中灰色)
        # 类别3 -> 255 (白色)
        
        gray_values = [0, 85, 170, 255]
        segmented_gray = np.zeros_like(mask, dtype=np.uint8)
        
        for i in range(min(len(gray_values), np.max(mask) + 1)):
            segmented_gray[mask == i] = gray_values[i]
        
        # 转换为PIL图像
        result_image = Image.fromarray(segmented_gray, mode='L')
        
        return result_image

def main():
    app = QApplication(sys.argv)
    window = ImageSegmentationGUI()
    window.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()